import numpy as np
import torchvision.datasets as datasets
import torch
import torch.nn as nn
from torchvision import transforms
from torch.utils.data.dataloader import DataLoader
import torch.optim as optim
import matplotlib.pylab as plt
import tqdm
plt.switch_backend('agg')
classes = ('plane','car','bird','cat','deer',
          'dog','frog','horse','ship','truck')

 # 设置transforms
transform = transforms.Compose([
    transforms.ToTensor(), # numpy -> Tensor
    transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))  # 归一化 ，范围[-1,1]
])
# 训练集,
'''
    下载数据集将download改为true
'''
trainset = datasets.CIFAR10(root='data/CIFAR10',train=True,download=False,transform=transform)
# 测试集
# testset = datasets.CIFAR10(root='./CIFAR10',train=False,download=False,transform=transform)
BATCH_SIZE = 2

train_loader = DataLoader(trainset,batch_size=BATCH_SIZE,shuffle=True,num_workers=0)
from model.convnext import *
model=convnext_tiny(pretrained=False,in_22k=False)
model=model.train()
model=model.cuda()
optimizer=optim.AdamW(model.parameters(),lr=0.0001)
Loss=nn.CrossEntropyLoss()
list1=[]
list2=[]
print(len(trainset))
if __name__=='__main__':
#
    for epoch in range(10):
        sumloss=0
        sum1=0
        numlen=0
        for (x,y) in tqdm.tqdm((train_loader)):
            numlen+=1
            x=x.cuda()
            y=y.cuda()
            y_pred=model(x)
            optimizer.zero_grad()
            yy=torch.argmax(y_pred,dim=1)
            num=sum(yy==y)
            sum1+=num
            loss=Loss(y_pred,y)
            sumloss+=loss.item()
            loss.backward()
            optimizer.step()
        acc=sum1/50000.0
        list1.append(sumloss/numlen)
        list2.append(acc)
        print(acc,sumloss/numlen)
    plt.plot(list1)
    plt.savefig('loss.jpg')
    plt.cla()
    plt.plot(list2)
    plt.savefig('acc.jpg')